import cv2
import numpy as np
import random
import os
import os.path as osp

CENTERS = np.array([[219.57673667, 220.39418417, 228.82714055],
                    [219.89855072, 181.42028986, 243.42028986]])


class Filter():
    def __init__(self):
        pass

    def predict(self, X: np.ndarray = None, centers: np.ndarray = CENTERS):
        if X is None:
            return None
        x = np.broadcast_to(X, centers.shape)
        y = np.sum((centers - x) ** 2, axis=1)
        return y.argmin(axis=0)

    def filter(self, image, cls):
        im = image.copy()
        contrast = im * 1.5
        contrast[contrast > 255] = 255
        contrast = contrast.astype(np.uint8)
        if cls == 0:
            lower_pink = np.array([150, 30, 30])

            upper_pink = np.array([179, 255, 255])
        else:
            lower_pink = np.array([150, 80, 80])

            upper_pink = np.array([179, 255, 255])
        hsv_image = cv2.cvtColor(contrast, cv2.COLOR_BGR2HSV)
        mask = cv2.inRange(hsv_image, lower_pink, upper_pink)
        closing = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((30, 30), np.uint8))
        # closing = cv2.morphologyEx(closing, cv2.MORPH_CLOSE, np.ones((10, 10), np.uint8))
        # closing = cv2.morphologyEx(closing, cv2.MORPH_CLOSE, np.ones((10, 10), np.uint8))
        opening = cv2.morphologyEx(closing, cv2.MORPH_OPEN, np.ones((10, 10), np.uint8))
        dilate = cv2.dilate(opening, np.ones((5, 5), dtype=np.uint8))
        # after_filted = cv2.bitwise_and(image, image, mask=opening)
        cnts, _ = cv2.findContours(dilate, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        return cnts, dilate


# for i in os.listdir(r'D:\PythonSpace\lianxi\pollen\tongzhou20190831'):
#     image = cv2.imread(osp.join(r'D:\PythonSpace\lianxi\pollen\tongzhou20190831', i))
#
#     classification = predict(np.expand_dims(np.percentile(image, 50, axis=(0, 1)), 0))
#     cnts, opening = filter(image, classification)
#     if len(cnts) > 0:
#         for j in cnts:
#             x, y, w, h = cv2.boundingRect(j)
#             if w * h < 500:
#                 continue
#             cv2.rectangle(image, (x - 10, y - 10), (x + w + 10, y + h + 10), (0, 255, 0), 1)
#     else:
#         cv2.imwrite(osp.join(r'D:\PythonSpace\lianxi\pollen\result\special', i), image)
#     cv2.imwrite(osp.join(r'D:\PythonSpace\lianxi\pollen\result\gary_detect', 'opening_' + i), opening)
#     cv2.imwrite(osp.join(r'D:\PythonSpace\lianxi\pollen\result\gary_detect', 'rgb_' + i), image)


# # contrast = image * 1.2
# # contrast[contrast > 255] = 255
# # contrast = contrast.astype(np.uint8)
# # image = cv2.imread(r'D:\PythonSpace\lianxi\pollen\pollen_data_multi_classes\JPEGImages\3_55.jpg')
# gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# # _, thre = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)
# # GBlur = cv2.GaussianBlur(gray, (5, 5), 0)
# # cv2.imshow('thre', thre)
# canny = cv2.Canny(gray, 50, 150)
# # closing = cv2.morphologyEx(canny, cv2.MORPH_CLOSE, np.ones((5,5), dtype='uint8'))
# contours, hierarchy = cv2.findContours(canny, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
#
# # im = image.copy()
# # cv2.drawContours(im, contours, -1, (0, 255, 0), 1)
# if len(contours) > 0:
#     new_cnts = []
#     for j in contours:
#         x, y, w, h = cv2.boundingRect(j)
#         if w * h < 800:
#             continue
#
#         cv2.rectangle(image, (x - 10, y - 10), (x + w + 10, y + h + 10), (0, 255, 0), 1)
#     # cv2.drawContours(image, np.array(new_cnts), -1, (0, 255, 0), 1)
# cv2.imwrite(osp.join(r'D:\PythonSpace\lianxi\pollen\result\gary_detect', 'rgb_' + i), image)
# cv2.imshow('im', im)
# cv2.imshow('image', image)
# cv2.imshow('im', im)
# cv2.imshow('contrast_gray', gray)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
# break

if __name__ == '__main__':
    a = np.array([[[0, 0, 0], [0, 0, 0]]])
    b = np.array([0, 0, 0])
    print()
    # f = Filter()
    # image = cv2.imread(r'/home/ubuntu/code/ssd.pytorch/data/pollen_data/train/JPEGImages/3_55.jpg')
    # classification = f.predict(np.expand_dims(np.percentile(image, 50, axis=(0, 1)), 0))
    # cnts, op = f.filter(image, classification)
    # # op = np.zeros((512, 512), dtype=np.uint8)
    # # op[100:300, 200:300] = 255
    # cv2.imshow('ss', op)
    #
    # for i in cnts:
    #     (x, y), radius = cv2.minEnclosingCircle(i)
    #     center, radius = (int(x), int(y)), int(radius)  # for the minimum enclosing circle
    #
    #     op = cv2.circle(op, center, radius+5, 255, -1)  # red
    # res = cv2.bitwise_and(image, image, mask=op)
    # # kernel = np.ones((9, 9), np.uint8)
    # # gradient = cv2.morphologyEx(op, cv2.MORPH_BLACKHAT, kernel)
    # # # a = cv2.erode(op,np.ones((3,3),np.uint8))
    # # #
    # cv2.imshow('a', res)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
